import paddle
from paddle import jit


class Test(paddle.nn.Layer):
    def __init__(self):
        super(Test, self).__init__()

    def forward(self, encodings, durations):
        batch_size, t_enc = paddle.shape(durations)
        slens = paddle.sum(durations, -1)
        t_dec = paddle.max(slens)
        t_dec_1 = t_dec + 1
        flatten_duration = paddle.cumsum(
            paddle.reshape(durations, [batch_size * t_enc])) + 1
        init = paddle.zeros(t_dec_1)
        m_batch = batch_size * t_enc
        M = paddle.zeros([t_dec_1, m_batch])
        for i in range(m_batch):
            d = flatten_duration[i]
            m = paddle.concat(
                [paddle.ones(d), paddle.zeros(t_dec_1 - d)], axis=0)
            M[:, i] = m - init
            init = m
        M = paddle.reshape(M, shape=[t_dec_1, batch_size, t_enc])
        M = M[1:t_dec_1, :, :]
        M = paddle.transpose(M, (1, 0, 2))
        encodings = paddle.matmul(M, encodings)
        return encodings


from paddle.static import InputSpec

if __name__ == '__main__':
    model = Test()
    encodings = paddle.ones((1, 10, 3), dtype=paddle.float32)
    durations = paddle.ones((1, 10), dtype=paddle.int32)
    model(encodings, durations)
    static_model = jit.to_static(model, [
        InputSpec((1, 10, 3), dtype=paddle.float32),
        InputSpec((1, 10), dtype=paddle.int32)
    ])
    jit.save(static_model, '/home/test')
